import os
from repository.SampleRepository import SampleRepository
from repository.FeatureRepository import FeatureRepository
from repository.HGMSimDO.SampleDO import SampleDO
from repository.HGMSimDO.FeatureDO import FeatureDO
import time
import subprocess
from basicTools.FileProcess import FileProcess
from datetime import datetime
import json
from PIL import Image
import io
from HGMSim.EasySample import EasySample


class SampleService:
    def __init__(self):
        """
        初始化文件
        """
        self.easySample = EasySample()
        self.sampleRepository = SampleRepository()
        self.featureRepository = FeatureRepository()

        with open("config.json", "r") as file:
            self.config = json.load(file)

        self.idaPath = self.config["featureExtract"]["idaPath"]
        self.scriptPath = self.config["featureExtract"]["scriptPath"]
        self.idaFeatureSaveDirectory = self.config["featureExtract"][
            "idaFeatureSaveDirectory"
        ]

        self.idaLogBasePath = self.config["featureExtract"]["idaLogBasePath"]
        if not os.path.exists(self.idaLogBasePath):
            os.makedirs(self.idaLogBasePath)

        self.idaLog = os.path.join(
            self.idaLogBasePath, self.config["featureExtract"]["idaLog"]
        )
        if not os.path.isfile(self.idaLog):
            with open(self.idaLog, "w") as f:
                pass

        self.sampleSaveDirectory = self.config["storage"]["sampleSaveDirectory"]
        self.md5SampleSaveDirectory = self.config["storage"]["md5SampleSaveDirectory"]
        self.greyImageDirectory = self.config["storage"]["greyImageDirectory"]

    def saveSample2DB(self, fileName, userId, userName, md5Value):
        sampleDO = SampleDO(
            None, userId, userName, datetime.now(), 0, fileName, md5Value
        )
        return self.sampleRepository.addSample(sampleDO)

    def savaFeature2DB(self, md5, famLabel, graphEmbedding, isTrainSample):
        featureDO = FeatureDO(None, md5, famLabel, graphEmbedding, isTrainSample)
        return self.addFeature2DB(featureDO)

    def extractFeature(self, filePath, sampleId, md5):
        """
        提取样本需要保存的特征，异步执行
        # todo 后续改成异步的实现方式
        """
        a = time.time()
        cmd = "{0} -L{1} -A -S{2} {3}".format(
            self.idaPath,
            self.idaLog,
            self.scriptPath,
            filePath,
        )
        print("执行命令" + cmd)
        subprocess.call(cmd)
        b = time.time()
        print("样本处理完成，耗时{}".format(b - a))

    def addFeature2DB(self, featrueDO):
        """
        将特征信息保存到数据库
        """
        return self.featureRepository.addFeature(featrueDO)

    def isPeFile(self, filePath):
        """
        判断对应路径下的文件是否为PE文件
        return true/false
        """
        fileProcess = FileProcess()
        return fileProcess.isPeFile(filePath)

    def removeSampleByPath(self, filePath):
        if os.path.exists(filePath):
            try:
                os.remove(filePath)
                return True
            except Exception as e:
                return False
        else:
            return False

    def saveOriginSample(self, file, savePath):
        """
        保存上传文件的字节流到本地
        return 保存是否成功
        """
        print("收到文件:" + file.filename)
        file.save(savePath)
        print("保存文件:" + file.filename)
        if self.isPeFile(savePath) == False:
            self.removeSampleByPath(savePath)
            print("删除文件:" + file.filename)
            return False
        return True

    def querySampleListByUser(self, userId):
        """
        查询用户上传的样本信息
        """
        return self.sampleRepository.querySampleListByUser(userId)

    def querySampleByMd5(self, name):
        """
        根据md5查询样本
        """
        return self.sampleRepository.querySampleListByMd5(name)

    def querySampleByFileName(self, name):
        """
        根据文件名查询样本
        """
        return self.sampleRepository.querySampleListByFileName(name)

    def queryFeatureByMd5(self, name):
        """
        通过md5值请求样本信息
        """
        return self.featureRepository.queryFeatureByMd5(name)

    def queryGreyImage(self, name):
        fileProcess = FileProcess()
        md5Value = fileProcess.calculateMd5(self.sampleSaveDirectory, name)
        img = Image.open(self.greyImageDirectory + md5Value + ".png")
        img_io = io.BytesIO()
        img.save(img_io, "PNG")
        img_io.seek(0)
        return img_io

    def querySampleSum(self):
        """
        查询样本总数
        """
        return self.sampleRepository.querySampleSum()

    def queryfamSum(self):
        """
        查询标签总数
        """
        return self.featureRepository.queryfamLableSum()

    def queryTodayUploadSampleSum(self):
        """
        查看今日上传的样本数量
        """
        return self.sampleRepository.queryTodayUploadSampleSum()

    def queryFamLabelTop10(self):
        """
        查看数量前10的样本标签
        """
        return self.featureRepository.queryFamLabelTop10()

    def calculateReuseFunction(self, name, name1):
        """
        获得两个样本中重复使用的函数
        """
        sample = self.easySample.queryProcessSample(name)
        sample1 = self.easySample.queryProcessSample(name1)
        return self.findAllSameFunction(sample, sample1)

    def findAllSameFunction(self, xData, yData):
        """
        用于判断两个样本中相同的函数有多少
        这里是粗略的判断，用嵌入值是否相同来确定
        """
        tempData = xData["att"]
        sample_x = [str(item) for item in tempData]
        hashmap_x = {}
        for i in range(len(sample_x)):
            att = sample_x[i]
            hashmap_x[att] = i

        tempData = yData["att"]
        sample_y = [str(item) for item in tempData]
        hashmap_y = {}
        for i in range(len(sample_y)):
            att = sample_y[i]
            hashmap_y[att] = i

        res = 0
        for att in hashmap_x.keys():
            if att in hashmap_y:
                res += 1

        xSum = 0
        ySum = 0
        # 统计以下边的个数
        tempData = xData["adj"]
        for row in tempData:
            for num in row:
                xSum += num
        tempData = yData["adj"]
        for row in tempData:
            for num in row:
                ySum += num
        ans = {}
        sampleX = {}
        sampleX["funcSize"] = len(sample_x)
        sampleX["edgeSize"] = xSum.item()
        ans["sampleX"] = sampleX
        sampleY = {}
        sampleY["funcSize"] = len(sample_y)
        sampleY["edgeSize"] = ySum.item()
        ans["sampleY"] = sampleY
        ans["reuseFunction"] = res
        # print(
        #     "样本x长度：{},边数量：{} 样本y长度：{}，边数量：{} 相同的函数个数：{}  占比：{}".format(
        #         len(sample_x),
        #         xSum,
        #         len(sample_y),
        #         ySum,
        #         res,
        #         round(res / len(sample_x), 2),
        #     )
        # )
        return ans
